package yycore

import chisel3._
import chisel3.util._
import chisel3.util.experimental.BoringUtils
import common.Constants._

class DecodeCtrlSignal() extends Bundle() {
  val instValid = Output(Bool())
  val pc = Output(UInt(32.W))
  val inst = Output(UInt(32.W))
  val op1_data = Output(UInt(DataBits.W))
  val rs2_data = Output(UInt(DataBits.W))
  val op2_data = Output(UInt(DataBits.W))
  val fu_typ = Output(FuType())
  val fu_fun = Output(UInt(fuOpType.W))
  val br_type = Output(BRType())
  val rf_wen = Output(Bool())
  val wb_addr = Output(UInt(5.W))
  val wb_sel = Output(UInt(WB_X.getWidth.W))
  val prepcNext = Output(UInt(AddrBits.W))
  val exceptionVec = Output(Vec(16, Bool()))
}

class DecodeStage extends Module with ExceptionNO{
  val io = IO(new Bundle() {
    val flush = Input(Bool())
    val in = Flipped(Decoupled(new InstrPackage))
    val out = Decoupled(new DecodeCtrlSignal())
    val r_reg = Flipped(new GPRReadIO)
    val bypass = Vec(2 ,new ByPassDataIO)
  })
  io := DontCare

  //val ld_hazard_stall = Wire(Bool())
  io.in.ready := io.out.ready //&& !ld_hazard_stall

  //**********************************
  // Pipeline State Registers
  // Instruction Decode State
  val reset_ifdec = WireInit(0.U.asTypeOf(new InstrPackage))
  reset_ifdec.inst := BUBBLE

  val dec_reg = RegInit(reset_ifdec)

  val inst = dec_reg.inst
  // **********************************
  // update state reg
  when(io.in.fire() && io.out.fire() && !io.flush) {
    dec_reg := io.in.bits
  }.elsewhen((!io.in.fire() && io.out.fire()) || io.flush) {
    dec_reg := reset_ifdec
  }

  // match instruction type to generate the control signal
  val csignals = ListLookup(inst, Control.decodeDefault, Control.DecodeTable)

  // put the control signals into variables
  val (cs_val_inst: Bool) :: cs_br_type :: cs_op1_sel :: cs_op2_sel :: cs0 = csignals
  val cs_rs1_ren :: cs_rs2_ren :: cs_fu_type :: cs_fu_fun :: cs_wb_sel :: (cs_rf_wen: Bool) :: Nil = cs0

  // fetch immediates
  val imm_i = inst(31, 20)
  //val imm_s = Cat(inst(31, 25), inst(11, 7))
  //val imm_b = Cat(inst(31), inst(7), inst(30, 25), inst(11, 8))
  val imm_u = inst(31, 12)
  //val imm_j = Cat(inst(31), inst(19, 12), inst(20), inst(30, 21))
  val imm_z = Cat(Fill(XLEN - 5, 0.U), inst(19,15))
  // output sign-extend immediates
  val imm_i_sext = Cat(Fill(XLEN - 12, imm_i(11)), imm_i)
  val imm_s_sext = Cat(Fill(XLEN - 12, inst(31)), Cat(inst(31, 25), inst(11, 7)))
  val imm_b_sext = Cat(Fill(XLEN - 13, inst(31)), Cat(inst(31), inst(7), inst(30, 25), inst(11, 8)), 0.U)
  val imm_u_sext = Cat(Fill(XLEN - 32, imm_u(19)), imm_u, Fill(12, 0.U))
  val imm_j_sext = Cat(Fill(XLEN - 21, inst(31)), Cat(inst(31), inst(19, 12), inst(20), inst(30, 21)), 0.U)

  val dec_op1_data = Wire(UInt(XLEN.W))
  val dec_op2_data = Wire(UInt(XLEN.W))
  val dec_rs2_data = Wire(UInt(XLEN.W))

  val dec_rs_addr = Wire(Vec(2, UInt(5.W)))
  dec_rs_addr(0) := inst(RS1_MSB, RS1_LSB)
  dec_rs_addr(1) := inst(RS2_MSB, RS2_LSB)
  val dec_wb_addr = inst(RD_MSB, RD_LSB)
  // read registers
  io.r_reg.rs1_addr := dec_rs_addr(0)
  io.r_reg.rs2_addr := dec_rs_addr(1)

  // bypass
  val rs1_data = MuxCase(io.r_reg.rs1_data, Array(
    (io.bypass(0).addr === dec_rs_addr(0) && (dec_rs_addr(0) =/= 0.U) && io.bypass(0).rf_wen)  -> io.bypass(0).data,
    (io.bypass(1).addr === dec_rs_addr(0) && (dec_rs_addr(0) =/= 0.U) && io.bypass(1).rf_wen)  -> io.bypass(1).data,
  ))
  val rs2_data = MuxCase(io.r_reg.rs2_data, Array(
    (io.bypass(0).addr === dec_rs_addr(1) && (dec_rs_addr(1) =/= 0.U) && io.bypass(0).rf_wen)  -> io.bypass(0).data,
    (io.bypass(1).addr === dec_rs_addr(1) && (dec_rs_addr(1) =/= 0.U) && io.bypass(1).rf_wen)  -> io.bypass(1).data,
  ))

  // set alu opcode data
  dec_op1_data := MuxCase(0.U, Array(
    (cs_op1_sel === OP1_RS1)  -> rs1_data,
    (cs_op1_sel === OP1_IMU)  -> imm_u_sext,
    (cs_op1_sel === OP1_IMJ)  -> imm_j_sext,
    (cs_op1_sel === OP1_IMZ)  -> imm_z,
  )).asUInt()
  dec_op2_data := MuxCase(0.U, Array(
    (cs_op2_sel === OP2_RS2)  -> rs2_data,
    (cs_op2_sel === OP2_IMI)  -> imm_i_sext,
    (cs_op2_sel === OP2_IMS)  -> imm_s_sext,
    (cs_op2_sel === OP2_IMB)  -> imm_b_sext,
    (cs_op2_sel === OP2_PC)   -> dec_reg.pc,
  )).asUInt()
  dec_rs2_data := rs2_data

  // ============= br ================
  val br_type = WireInit(cs_br_type)

  when (BRType.isBranch(cs_br_type)) {
    def isLink(reg: UInt) = (reg === 1.U || reg === 5.U)
    when (isLink(dec_wb_addr) && cs_br_type === BRType.jal) { br_type := BR_CALL }
    when (cs_br_type === BRType.jalr) {
      when (isLink(dec_rs_addr(0))) { br_type := BR_RET }
      when (isLink(dec_wb_addr)) { br_type := BR_CALL }
    }
  }

  // ================ CSR ================
  val isCSR = cs_fu_type === FuType.csr
  val csr_ren = isCSR && (cs_fu_fun === CSROpType.S || cs_fu_fun === CSROpType.C) && (cs_rs1_ren.asBool() && dec_rs_addr(0) === 0.U)
  val fu_fun = Mux(csr_ren, CSROpType.R, cs_fu_fun)

  // interrupt epc
//  BoringUtils.addSource(dec_reg.pc, "dec_pc")
//  BoringUtils.addSource(dec_reg.instValid, "dec_instValid")

  // **********************************
  // next ctrl signals
  io.out.valid := true.B
  io.out.bits.instValid := dec_reg.instValid
  io.out.bits.br_type := br_type
  io.out.bits.rf_wen := cs_rf_wen
  io.out.bits.fu_typ := cs_fu_type
  io.out.bits.fu_fun := fu_fun
  io.out.bits.wb_sel := cs_wb_sel

  io.out.bits.pc := dec_reg.pc
  io.out.bits.inst := inst
  io.out.bits.op1_data := dec_op1_data
  io.out.bits.op2_data := dec_op2_data
  io.out.bits.rs2_data := dec_rs2_data
  io.out.bits.wb_addr := dec_wb_addr
  //io.out.bits.csr_cmd := Mux(csr_ren, CSROpType.R, cs_csr_cmd)

  io.out.bits.prepcNext := dec_reg.prepcNext
  io.out.bits.exceptionVec.map(_ := false.B)
  io.out.bits.exceptionVec(instrPageFault) := cs_val_inst
}
